#version 320 es
//输入布局限定符声明一个16 X 16(Xl)
//的本地工作组
layout (local_size_x = 16, local_size_y = 16) in;

layout(location = 0) uniform highp sampler2D WkN;    //fft 正交旋转因子数据
layout(location = 1) uniform lowp isampler2D mirrorByteLUT; //字节镜像查找表

layout(location = 20) uniform lowp uint iscoldir;   //是否按列进行变换,0:每行数据进行变换 1:每列数据进行变换
layout(location = 21) uniform mediump int bitscount;   //进行变换的数据长度转为 2^N 的N-1值,这个值也是要进行的蝶形计算次数
layout(location = 22) uniform lowp int calcbufStarti;   //计算缓冲数组的索引起始值
layout(location = 23) uniform mediump uvec2 bufsize;   //用于计算的缓冲区尺寸
layout(location = 24) uniform highp float alphaImag;     //在WkN虚部前相乘的数，可用于逆变换（=-1）
layout(location = 25) uniform highp float mulConst;     //蝶形变换后除以的常数，可用于使浮点数变小，以提高精度
layout(location = 26) uniform mediump int cyc;          //循环执行本计算着色器的次数

layout (binding = 0, r32f) highp uniform image2DArray calcbuf;  //用于计算的缓冲区数组 改为外循环后 不需要coherent了 因为读写都是不同的显存

void main(void)
{
    //cyc等于-1表明进行输入数据重排序
    if( cyc == -1 && gl_GlobalInvocationID.x < bufsize.x && gl_GlobalInvocationID.y < bufsize.y ){
        //首先进行抽序
        highp vec2 complex;
        mediump ivec2 srcpos;
        srcpos.y = int(iscoldir == 0u) * int(gl_GlobalInvocationID.x)
                + int(iscoldir == 1u) * int(gl_GlobalInvocationID.y);
        srcpos.x = 0;
        //只进行3次，也就是最大是24位
        srcpos.x = texelFetch(mirrorByteLUT,ivec2(srcpos.y & 0x0FF,0),0).x;  //镜像索引得到要抽取的位置
        srcpos.x = srcpos.x << 8;
        srcpos.y = srcpos.y >> 8;
        srcpos.x = srcpos.x | texelFetch(mirrorByteLUT,ivec2(srcpos.y & 0x0FF,0),0).x;
        srcpos.x = srcpos.x << 8;
        srcpos.y = srcpos.y >> 8;
        srcpos.x = srcpos.x | texelFetch(mirrorByteLUT,ivec2(srcpos.y & 0x0FF,0),0).x;
        srcpos.x = srcpos.x >> (24 - bitscount); //只要最左面的bitscount位数

        srcpos = int(iscoldir == 0u) * ivec2(srcpos.x,gl_GlobalInvocationID.y)
                + int(iscoldir == 1u) * ivec2(gl_GlobalInvocationID.x,srcpos.x);
        srcpos.x = srcpos.x * 2;

        //读取抽取位置的值 写入缓冲区
        complex.x = imageLoad(calcbuf,ivec3(srcpos,calcbufStarti)).x;   //读实部
        srcpos.x = srcpos.x + 1;
        complex.y = imageLoad(calcbuf,ivec3(srcpos,calcbufStarti)).x;    //读虚部
        complex = mulConst * complex;   //总体乘以一个倍率
        imageStore(calcbuf,
                   ivec3(2u * gl_GlobalInvocationID.x,gl_GlobalInvocationID.y,1-calcbufStarti),
                   vec4(complex.x,0,0,0));  //写实部
        imageStore(calcbuf,
                   ivec3(1u + 2u * gl_GlobalInvocationID.x,gl_GlobalInvocationID.y,1-calcbufStarti),
                   vec4(complex.y,0,0,0)); //写虚部
        return;
    }

    if( cyc >= 0 && gl_GlobalInvocationID.x < bufsize.x && gl_GlobalInvocationID.y < bufsize.y ){
        highp vec4 complex;
        mediump ivec4 srcpos;
        srcpos.x = int(iscoldir == 0u) * int(gl_GlobalInvocationID.x)
                + int(iscoldir == 1u) * int(gl_GlobalInvocationID.y);
        srcpos.y = srcpos.x % (1<<(cyc+1));   //组内序号
        srcpos.z = srcpos.y / (1<<cyc);   //分成上下半组，0上半组，1下半组
        srcpos.w = srcpos.y % (1<<cyc);   //旋转WkN序号索引
        srcpos.w = srcpos.w * (1<<(bitscount-1-cyc));
        //查找旋转因子WkN
        complex.zw = texelFetch(WkN, ivec2(srcpos.w,0), 0).xy;
        complex.zw = float(srcpos.z==0)*complex.zw - 1.0f * float(srcpos.z==1)*complex.zw; //下半组取负数
        complex.w = complex.w * alphaImag;  //虚部正负
        //需要乘以旋转因子的索引
        srcpos.x = int(srcpos.z==1)*srcpos.x + int(srcpos.z==0)*(srcpos.x + (1<<cyc));
        srcpos.xy = int(iscoldir == 0u) * ivec2(srcpos.x,gl_GlobalInvocationID.y)
                + int(iscoldir == 1u) * ivec2(gl_GlobalInvocationID.x,srcpos.x);
        srcpos.x = srcpos.x * 2;
        //读取蝶形变换左面数据
        complex.x = imageLoad(calcbuf,ivec3(srcpos.xy,calcbufStarti)).x;   //读实部
        srcpos.x = srcpos.x + 1;
        complex.y = imageLoad(calcbuf,ivec3(srcpos.xy,calcbufStarti)).x;    //读虚部
        //乘以旋转因子
        complex.zw = vec2(complex.z * complex.x - complex.w * complex.y,
                          complex.z * complex.y + complex.w * complex.x
                          );    //复数乘法
        //不用乘以旋转因子的索引
        srcpos.x = int(iscoldir == 0u) * int(gl_GlobalInvocationID.x)
                + int(iscoldir == 1u) * int(gl_GlobalInvocationID.y);
        srcpos.x = int(srcpos.z==0)*srcpos.x + int(srcpos.z==1)*(srcpos.x - (1<<cyc));
        srcpos.xy = int(iscoldir == 0u) * ivec2(srcpos.x,gl_GlobalInvocationID.y)
                + int(iscoldir == 1u) * ivec2(gl_GlobalInvocationID.x,srcpos.x);
        srcpos.x = srcpos.x * 2;
        //读取蝶形变换左面数据
        complex.x = imageLoad(calcbuf,ivec3(srcpos.xy,calcbufStarti)).x;   //读实部
        srcpos.x = srcpos.x + 1;
        complex.y = imageLoad(calcbuf,ivec3(srcpos.xy,calcbufStarti)).x;    //读虚部
        //相加得出结果
        complex.x = complex.x + complex.z;
        complex.y = complex.y + complex.w;
        //写入结果
        imageStore(calcbuf,
                   ivec3(2u * gl_GlobalInvocationID.x,gl_GlobalInvocationID.y,1-calcbufStarti),
                   vec4(complex.x,0,0,0));  //写实部
        imageStore(calcbuf,
                   ivec3(1u + 2u * gl_GlobalInvocationID.x,gl_GlobalInvocationID.y,1-calcbufStarti),
                   vec4(complex.y,0,0,0)); //写虚部
    }

}
